#!/usr/bin/env python3  # 指定用系统里的 python3 解释器运行本脚本
# -*- coding: utf-8 -*-  # 声明源文件使用 UTF-8 编码，避免中文注释在某些环境下乱码

import os  # 操作系统相关功能：读取环境变量、路径处理等
import sys  # 与解释器交互：标准输入输出、退出等
import argparse  # 命令行参数解析库
import signal  # 捕获 Ctrl+C 等信号，做优雅退出
from typing import List, Dict, Any  # 类型注解：便于阅读与 IDE 提示

# 允许用 .env 方式配置 DEEPSEEK_API_KEY（可选，不依赖则忽略异常）
try:
    from dotenv import load_dotenv  # 从 python-dotenv 导入加载函数
    load_dotenv()  # 自动读取当前目录或上层目录的 .env 文件，将键值对注入到环境变量
except Exception:
    pass  # 如果未安装或加载失败，不中断程序（仍可用系统环境变量或 CLI 参数）

# OpenAI 兼容 SDK（DeepSeek 提供 OpenAI 兼容的 API 端点）
try:
    from openai import OpenAI  # 客户端类：用于创建会话并发起请求
except ImportError:
    print("未安装 openai 库，请先执行: pip install -U openai", file=sys.stderr)  # 友好提示
    sys.exit(1)  # 依赖缺失无法继续，非零码退出


# 交互式 CLI 的横幅文案，用于提示可用的快捷命令
BANNER = """\
DeepSeek CLI (stream)
快捷命令: /reset 清空上下文 | /model 切换模型 | /exit 退出
"""


def parse_args():  # 封装命令行参数解析，便于 main 调用
    p = argparse.ArgumentParser(
        description="DeepSeek 命令行实时对话（流式打印，带会话记忆)"  # 帮助信息
    )
    p.add_argument(
        "--model",
        default="deepseek-chat",  # 默认模型：非思考模式（成本低、延迟小）
        help="模型名，如 deepseek-chat 或 deepseek-reasoner（V3.1 思考模式）",  # 使用说明
    )
    p.add_argument(
        "--system",
        default="You are a helpful assistant.",  # 默认系统提示词
        help="系统提示词（system prompt）",  # 自定义助手语气/行为
    )
    p.add_argument(
        "--max-tokens",
        type=int,
        default=None,  # 不指定时使用服务端默认
        help="生成上限 tokens。不填使用平台默认。",
    )
    p.add_argument(
        "--show-cot",
        action="store_true",  # 出现该开关则为 True，否则为 False
        help="使用 deepseek-reasoner 时，同时实时显示 reasoning_content（推理流）",  # 仅对 reasoner 有意义
    )
    p.add_argument(
        "--no-color",
        action="store_true",  # 禁用 ANSI 颜色输出
        help="禁用控制台颜色",
    )
    return p.parse_args()  # 解析并返回命令行参数对象


def colorize(s: str, code: str, enable: bool) -> str:  # 简单的彩色输出封装
    if not enable or not sys.stdout.isatty():  # 若禁用或非交互终端，则直接返回原字符串
        return s
    return f"\033[{code}m{s}\033[0m"  # 拼接 ANSI 转义序列，应用颜色/样式


def main():  # 主流程：初始化 -> 交互循环
    args = parse_args()  # 读取命令行参数

    api_key = os.getenv("DEEPSEEK_API_KEY")  # 从环境变量读取 API Key（也可通过 .env 注入）
    if not api_key:  # 若未设置则无法鉴权
        print("环境变量 DEEPSEEK_API_KEY 未设置。请先设置 API Key。", file=sys.stderr)
        sys.exit(2)  # 使用约定的非零退出码指示配置错误

    # DeepSeek 的 OpenAI 兼容端点：用 OpenAI SDK 直连 DeepSeek
    client = OpenAI(api_key=api_key, base_url="https://api.deepseek.com")  # 创建客户端实例

    # 会话消息历史；注意：服务端无状态，需在本地维护完整 messages 列表
    messages: List[Dict[str, Any]] = []  # 初始化空消息历史
    if args.system:  # 若提供了 system 提示词
        messages.append({"role": "system", "content": args.system})  # 设定对话基调

    # Ctrl+C（SIGINT）优雅退出：避免出现半行输出或脏状态
    def sigint_handler(signum, frame):  # 定义信号处理函数
        print("\n已退出。")  # 给用户反馈
        sys.exit(0)  # 正常退出
    signal.signal(signal.SIGINT, sigint_handler)  # 绑定 SIGINT 到处理函数

    print(BANNER)  # 打印横幅与快捷命令提示

    current_model = args.model  # 当前使用的模型名（可在会话中动态切换）
    print(f"当前模型: {colorize(current_model, '36', not args.no_color)}")  # 青色高亮显示

    while True:  # 进入 REPL 循环
        try:
            user = input(colorize("\n你 > ", "32;1", not args.no_color)).strip()  # 读取用户输入并去除首尾空白
        except EOFError:  # 捕获 Ctrl+Z / 管道结束等情况
            print("\n已退出。")
            break  # 跳出循环，结束程序

        if not user:  # 空输入直接下一轮
            continue

        # 处理快捷命令：/exit 退出、/reset 清空上下文、/model 切换模型
        if user == "/exit":  # 用户显式退出
            print("已退出。")
            break
        if user == "/reset":  # 清空历史（保留 system）
            messages = []  # 重置消息列表
            if args.system:
                messages.append({"role": "system", "content": args.system})  # 重新放入 system 提示
            print(colorize("（已清空上下文）", "33", not args.no_color))  # 黄色提示
            continue  # 进入下一轮输入
        if user.startswith("/model"):  # 动态切换模型：/model deepseek-reasoner
            parts = user.split(maxsplit=1)  # 只切一刀，保留模型名整体
            if len(parts) == 2:
                current_model = parts[1].strip()  # 更新当前模型名
                print(f"已切换模型为: {colorize(current_model, '36', not args.no_color)}")
            else:
                print("用法: /model deepseek-chat | deepseek-reasoner | 其它可用模型名")  # 使用帮助
            continue  # 本轮不与模型对话

        # 将用户输入加入消息历史，构成多轮上下文
        messages.append({"role": "user", "content": user})  # 附带最新 user 消息

        # 构造 Chat Completions 参数；开启流式，以边到边出
        create_kwargs = {
            "model": current_model,  # 本轮使用的模型
            "messages": messages,  # 完整历史（system + 若干轮 user/assistant）
            "stream": True,  # 启用流式响应
        }
        if args.max_tokens is not None:  # 可选的生成上限
            create_kwargs["max_tokens"] = args.max_tokens

        print(colorize("助理 >", "34;1", not args.no_color), end="", flush=True)  # 蓝色前缀，不换行

        # 对 deepseek-reasoner：增量片段将分为 reasoning_content（思考流）与 content（最终答案）
        # 注意：下一轮历史只保留最终 content，切勿把 reasoning_content 写回，否则会导致请求错误
        # 参考文档：https://api.deepseek.com （官方指南）
        final_answer_chunks = []  # 收集最终答案的增量片段
        cot_chunks = []  # 可选：收集思考流（仅展示用，不写回历史）

        try:
            stream = client.chat.completions.create(**create_kwargs)  # 发起流式对话请求，返回增量迭代器

            for chunk in stream:  # 逐块读取服务端推送的增量
                # 增量数据统一在 choices[0].delta；做一层防御式获取
                delta = None
                try:
                    delta = chunk.choices[0].delta  # 兼容 OpenAI 风格的字段
                except Exception:
                    delta = None  # 某些异常片段（如心跳）可忽略

                # 思考流（仅 reasoner 模型会提供该字段）
                piece_cot = getattr(delta, "reasoning_content", None) if delta else None  # 取出推理内容
                if piece_cot:
                    cot_chunks.append(piece_cot)  # 缓存以便需要时整体查看
                    if args.show_cot:  # 若用户要求展示思考流
                        print(colorize(piece_cot, "35", not args.no_color), end="", flush=True)  # 洋红色打印
                    continue  # 思考流不混入最终答案行

                # 正式内容（给用户看的最终文本）
                piece = getattr(delta, "content", None) if delta else None  # 取出本次的文本增量
                if piece:
                    final_answer_chunks.append(piece)  # 缓存
                    print(piece, end="", flush=True)  # 直接在终端流式打印

            print()  # 本轮输出结束后换行

        except KeyboardInterrupt:  # 用户在生成中按下 Ctrl+C
            print(colorize("\n[中断本轮生成]", "33", not args.no_color))  # 提示已中断
            messages.pop()  # 回滚刚刚 append 的 user，避免半条问题污染上下文
            continue  # 进入下一轮输入
        except Exception as e:  # 其它异常（网络、鉴权、配额等）
            print(colorize(f"\n[出错] {e}", "31", not args.no_color))  # 红色错误提示
            # 同样回滚本轮 user，保持历史干净
            if messages and messages[-1]["role"] == "user" and messages[-1]["content"] == user:
                messages.pop()
            continue  # 允许继续下一轮

        # 将本轮最终答案（拼接后的 content）写回历史，供下一轮上下文使用
        assistant_text = "".join(final_answer_chunks).strip()  # 合并增量并去除首尾空白
        messages.append({"role": "assistant", "content": assistant_text})  # 仅保存最终答案，不保存 COT


if __name__ == "__main__":  # Python 脚本入口惯用写法
    main()  # 调用主函数
